# %% [markdown]
# ### Import Libraries And Set Para

# %%
# import required libraries
import numpy as np
from datetime import datetime

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader

from torchvision import datasets, transforms

import matplotlib.pyplot as plt

# %%
# set parameters
LEARNING_RATE = 0.001
BATCH_SIZE = 32
EPOCHS = 20
RANDOM_SEED = 20

IMG_SIZE = 32
CLASSES_NUM = 10

# %% [markdown]
# ### Create datasets

# %%
# Define transforms. The input size of Lenet-5 is 32*32
transforms = transforms.Compose([transforms.Resize((32, 32)),
                                 transforms.ToTensor()])

# download and create datasets
train_dataset = datasets.MNIST(root='mnist_data',
                               train=True,
                               transform=transforms,
                               download=True)

valid_dataset = datasets.MNIST(root='mnist_data',
                               train=False,
                               transform=transforms)

# Define the data loaders
train_loader = DataLoader(dataset=train_dataset,
                          batch_size=BATCH_SIZE,
                          shuffle=True)

valid_loader = DataLoader(dataset=valid_dataset,
                          batch_size=BATCH_SIZE,
                          shuffle=False)

# %% [markdown]
# ### Define LeNet-5 network

# %%
# Define the LeNet-5 network
class LenNet5(nn.Module):
    def __init__(self, n_classes):
        super(LenNet5, self).__init__()

        # inchannels, out_channels, kernel_size, stride=1
        self.feature_extractor = nn.Sequential(            
            nn.Conv2d(1, 6, 5),
            nn.Tanh(),
            nn.AvgPool2d(kernel_size=2),
            nn.Conv2d(6, 16, 5),
            nn.Tanh(),
            nn.AvgPool2d(kernel_size=2),
            nn.Conv2d(16, 120, 5),
            nn.Tanh()
        )

        # two affine operations: y = Wx + b
        self.classifier = nn.Sequential(
            nn.Linear(in_features=120, out_features=84),
            nn.Tanh(),
            nn.Linear(in_features=84, out_features=n_classes),
        )

    def forward(self, x):
        x = self.feature_extractor(x)
        x = torch.flatten(x, 1)
        logits = self.classifier(x)
        probs = F.softmax(logits, dim=1)
        return logits, probs


# %% [markdown]
# ### Define functions to aid analysis

# %%
def plot_losses(train_losses, valid_losses):
    '''
    Function for plotting trainning and validation lossed
    '''

    plt.style.use('seaborn')

    train_losses = np.array(train_losses)
    valid_losses = np.array(valid_losses)

    fig, ax = plt.subplots(figsize=(8, 4.5))
    ax.plot(train_losses, color='blue', label='Training loss')
    ax.plot(valid_losses, color='red', label='Validation loss')
    ax.set(title='Loss over epochs', xlabel='Epoch', ylabel='Loss')
    ax.legend()
    fig.show()

    plt.style.use('default')
    return None

# %%
def get_accuracy(model, data_loader):
    '''
    Function for computing the accuracy of the predictions over the entire data_loader
    '''
    
    correct_pred = 0 
    n = 0
    
    with torch.no_grad():
        model.eval()
        for X, y_true in data_loader:

            _, y_prob = model(X)
            _, predicted_labels = torch.max(y_prob, 1)

            n += y_true.size(0)
            correct_pred += (predicted_labels == y_true).sum()

    return correct_pred.float() / n

# %% [markdown]
# ### Define Function used for training

# %%
# Define helpful functions used for training
def train(train_loader, model, criterion, optimizer):
    '''
    Function for the training step of the training loop
    '''
    model.train()
    runnin_loss = 0

    for X, y_true in train_loader:
        optimizer.zero_grad()

        # Forward: getting the predictions for the batch using current weight
        y_hat, _ = model(X)
        # Calculate the value of the loss function
        loss = criterion(y_hat, y_true)
        runnin_loss += loss.item() * X.size(0)

        # Backward: The weights are adjusted based on the loss
        loss.backward()
        optimizer.step()
        
    epoch_loss = runnin_loss / len(train_loader.dataset)
    return model, optimizer, epoch_loss

# %%
# define the function responsible for validation
def validate(valid_loader, model, criterion):
    '''
    Function for the validation step of the training loop
    '''

    model.eval()
    running_loss = 0

    for X, y_true in valid_loader:
        # Forward pass and record loss
        y_hat, _ = model(X)
        loss = criterion(y_hat, y_true)
        running_loss += loss.item() * X.size(0)

    epoch_loss = running_loss / len(valid_loader.dataset)
    return model, epoch_loss

# %%
# Define training loop
def training_loop(model, criterion, optimizer, train_loader, valid_loader, epochs, print_every=1):
    '''
    Function defining the entire training loop.

    '''

    # set objects for storing metrics
    best_loss = 1e10
    train_losses = []
    valid_losses = []

    # Train model
    for epoch in range(0, epochs):
        # training
        model, optimizer, train_loss = train(train_loader, model, criterion, optimizer)
        train_losses.append(train_loss)

        with torch.no_grad():
            model, valid_loss = validate(valid_loader, model, criterion)
            valid_losses.append(valid_loss)
        
        if (epoch % print_every) == (print_every - 1):

            # Calculate the accuracy of the model
            train_acc = get_accuracy(model, train_loader)
            valid_acc = get_accuracy(model, valid_loader)
            print(f'{datetime.now().time().replace(microsecond=0)} --- '
                  f'Epoch: {epoch}\t'
                  f'Train loss: {train_loss:.4f}\t'
                  f'Valid loss: {valid_loss:.4f}\t'
                  f'Train accuracy: {100 * train_acc:.2f}\t'
                  f'Valid accuracy: {100 * valid_acc:.2f}')
    
    plot_losses(train_losses, valid_losses)

    return model, optimizer, (train_losses, valid_losses)

# %% [markdown]
# ### Instantiate the model and train the net

# %%
torch.manual_seed(RANDOM_SEED)

model = LenNet5(CLASSES_NUM)
optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)
criterion = nn.CrossEntropyLoss()

# Start training
model, optimizer, _ = training_loop(model, criterion, optimizer, train_loader, valid_loader, EPOCHS)

# %%
ROW_IMG = 15
N_ROWS = 4

fig = plt.figure()
for index in range(1, ROW_IMG * N_ROWS + 1):
    plt.subplot(N_ROWS, ROW_IMG, index)
    plt.axis('off')
    plt.imshow(valid_dataset.data[index], cmap='gray_r')
    
    with torch.no_grad():
        model.eval()
        _, probs = model(valid_dataset[index][0].unsqueeze(0))
        
    title = f'{torch.argmax(probs)} ({torch.max(probs * 100):.0f}%)'
    
    plt.title(title, fontsize=7)
fig.suptitle('LeNet-5 - predictions');


